from lazyllm import (
     OnlineChatModule,OnlineEmbeddingModule,Retriever,Document,ChatPrompter,
     fc_register,ReactAgent
)


# base_url_http = "https://dashscope.aliyuncs.com/compatible-mode"
# llm=OnlineChatModule(source='qwen', model="qwen-plus", stream=False)

# while True:
#      query = input("query(enter 'quit' to exit): ")
#      if query == "quit":
#          break
#      res = llm.forward(query)
#      print(f"answer: {res}") 






# base_url_http = "https://dashscope.aliyuncs.com/compatible-mode"
# doc_path="F:\\file\\my\\workplace\\agentLazyLLM\\lazyllm-venv\\docs"


# embed_model = OnlineEmbeddingModule(source='qwen', embed_model_name='text-embedding-v4',type="embed")
# doc = Document(dataset_path=doc_path, embed=embed_model)
# retriever = Retriever(doc, group_name='CoarseChunk', similarity="bm25_chinese", topk=3)
# retriever_result = retriever(query="管理员性别")      

# context_str = "|".join([node.get_content() for node in retriever_result])
# print("context_str:"+context_str)


base_url_http = "https://dashscope.aliyuncs.com/compatible-mode"
doc_path="F:\\file\\my\\workplace\\agentLazyLLM\\lazyllm-venv\\docs"
embed_model = OnlineEmbeddingModule(source='qwen', embed_model_name='text-embedding-v4',type="embed")
doc = Document(dataset_path=doc_path, embed=embed_model)
retriever = Retriever(doc, group_name='CoarseChunk', similarity="bm25_chinese", topk=3)
llm=OnlineChatModule(source='qwen', model="qwen-plus", stream=False,base_url=base_url_http)
# prompt 设计
prompt = '你是一个AI,说话风趣幽默，有耐心'
llm.prompt(ChatPrompter(instruction=prompt, extra_keys=['context_str']))

@fc_register("tool")
def search_knowledge_base(query:str):
    """
    搜索知识库并返回相关文档内容
    
    Args:
        query (str): 搜索查询字符串
    """
    print(query)
    doc_node_list = retriever(query=query)
    context_str = "".join([node.get_content() for node in doc_node_list])
    return context_str

agent = ReactAgent(
    llm,
    tools=['search_knowledge_base'],
    prompt=prompt,
    stream=False
)
res = agent.forward("管理员性别")
print(f"answer: {res}") 